# -*- coding:utf-8 -*-
'''
@paper:A New Adaptive Weighted Mean Filter for Removing Salt-and-Pepper Noise
'''
import numpy as np

# Adaptive Weighted Mean Filter
class Awmf(object):
    def __init__(self, h, w_max):
        self.h = h
        self.w_max = w_max

    def t_zero_pad_matrix(self, A, t):
        '''在图像四周填充t层0'''
        m, n = np.shape(A)
        img_pad = np.zeros((m + t, n + t))
        img_pad[t:, t:] = A
        return img_pad

    def k_approximate_matrix(self, img_pad, i, j, k, t):
        k_approximate_mat = img_pad[(i + t - k):(i + t + k + 1), (j + t - k):(j + t + k + 1)]
        return k_approximate_mat

    def row_matrix(self, img_k_approximate):
        H = img_k_approximate.flatten()
        H = np.sort(H)
        return H

    def amf(self, img):
        '''Adaptive Mean Filter'''
        m, n = np.shape(img)
        z = img.copy()
        img_pad = self.t_zero_pad_matrix(img, self.w_max)
        for i in range(m):
            for j in range(n):
                w = 1
                while w <= self.w_max:
                    k_approximate_mat = self.k_approximate_matrix(img_pad, i, j, k=w, t=self.w_max)
                    median_value = np.median(k_approximate_mat)
                    max_value = np.max(k_approximate_mat)
                    min_value = np.min(k_approximate_mat)
                    if min_value < median_value < max_value:
                        if min_value < img[i, j] < max_value:
                            z[i, j] = img[i, j]
                            break
                        else:
                            z[i, j] = median_value
                            break
                    else:
                        w = w + self.h
                        if w > self.w_max:
                            z[i, j] = median_value
        return z

    def mean_awmf(self, Y, max_value, min_value):
        m, n = np.shape(Y)
        A = np.zeros(np.shape(Y))
        for k in range(m):
            for l in range(n):
                if min_value < Y[k, l] < max_value:
                    A[k, l] = 1
        if np.sum(A) != 0:
            mean_value = np.sum(A * Y) / np.sum(A)
        else:
            mean_value = -1
        return mean_value

    def process_image(self, img):
        m, n = np.shape(img)
        z = img.copy()
        img_pad = self.t_zero_pad_matrix(img, self.w_max)
        for i in range(m):
            for j in range(n):
                w = 1
                while w <= self.w_max:
                    k_approximate_mat1 = self.k_approximate_matrix(img_pad, i, j, k=w, t=self.w_max)
                    k_approximate_mat2 = self.k_approximate_matrix(img_pad, i, j, k=w + self.h, t=self.w_max)
                    max_value1 = np.max(k_approximate_mat1)
                    min_value1 = np.min(k_approximate_mat1)
                    mean_value1 = self.mean_awmf(k_approximate_mat1, max_value1, min_value1)
                    max_value2 = np.max(k_approximate_mat2)
                    min_value2 = np.min(k_approximate_mat2)
                    if min_value1 == min_value2 and max_value1 == max_value2 and mean_value1 != -1:
                        if min_value1 < img[i, j] < max_value1:
                            z[i, j] = img[i, j]
                            break
                        else:
                            z[i, j] = mean_value1
                            break
                    else:
                        w = w + self.h
                        if w > self.w_max:
                            z[i, j] = mean_value1
        return z
